package com.zhangxin.aiInterview.manager.wsserver;

import com.zhangxin.aiInterview.common.ErrorCode;
import com.zhangxin.aiInterview.exception.BusinessException;
import org.springframework.web.socket.WebSocketSession;

import java.io.IOException;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 前端ws连接管理
 */
public class WsManager {
    // Key: 用户ID, Value: 单个WebSocket连接
    private static final ConcurrentHashMap<Long, WebSocketSession> wsConnections = new ConcurrentHashMap<>();

    /**
     * 添加或更新连接
     * @param userId 用户唯一标识
     * @param ws WebSocket连接对象
     */
    public static void addOrUpdateConnection(Long userId, WebSocketSession ws) {
        // 先关闭旧连接（避免泄漏）
        WebSocketSession oldWs = wsConnections.get(userId);
        if (oldWs != null) {
            try {
                oldWs.close();
            } catch (IOException e) {
                throw new BusinessException(ErrorCode.OPERATION_ERROR,"关闭ws连接失败");
            }
        }
        wsConnections.put(userId, ws);
    }

    /**
     * 获取连接
     * @param userId
     * @return
     */
    public static WebSocketSession getConnection(Long userId) {
        return wsConnections.get(userId);
    }

    /**
     * 移除连接
     * @param userId 用户ID
     */
    public static void removeConnection(Long userId) {
        WebSocketSession ws = wsConnections.get(userId);
        if (ws != null) {
            try {
                ws.close();
            } catch (IOException e) {
                throw new BusinessException(ErrorCode.OPERATION_ERROR,"关闭ws连接失败");
            }
        }
        wsConnections.remove(userId);
    }
}